#include "kalman_filter.hpp"

KalmanFilter::KalmanFilter(double R_value, double Q_value)
    : x(4, 1), x_ex(4, 1), F(4, 4), H(2, 4), P(4, 4), R(2, 2), Q(4, 4)
{
    x = Matrix<double>::Zero(4, 1);
    x_ex = Matrix<double>::Zero(4, 1);
    F = Matrix<double>::Identity(4);
    
    H(0, 0) = 1; H(0, 1) = 0; H(0, 2) = 0; H(0, 3) = 0;
    H(1, 0) = 0; H(1, 1) = 1; H(1, 2) = 0; H(1, 3) = 0;
    
    P = Matrix<double>::Identity(4) * 1000.0;
    R = Matrix<double>::Identity(2) * R_value;
    Q = Matrix<double>::Identity(4) * Q_value;
}

void KalmanFilter::update_state_transition_matrix(double dt) {
    F(0, 0) = 1; F(0, 1) = 0; F(0, 2) = dt; F(0, 3) = 0;
    F(1, 0) = 0; F(1, 1) = 1; F(1, 2) = 0;  F(1, 3) = dt;
    F(2, 0) = 0; F(2, 1) = 0; F(2, 2) = 1;  F(2, 3) = 0;
    F(3, 0) = 0; F(3, 1) = 0; F(3, 2) = 0;  F(3, 3) = 1;
}

void KalmanFilter::predict() {
    if (!x.is_valid() || !F.is_valid() || !P.is_valid() || !Q.is_valid()) {
        throw std::runtime_error("Invalid matrix state in predict");
    }
    x = F * x;
    P = F * P * F.transpose() + Q;
}

void KalmanFilter::extrapolate() {
    if (!x.is_valid() || !F.is_valid()) {
        throw std::runtime_error("Invalid matrix state in extrapolate");
    }
    x_ex = F * x;
}

void KalmanFilter::update(const Matrix<double>& z) {
    if (z.get_rows() != 2 || z.get_cols() != 1) {
        throw std::invalid_argument("Invalid measurement vector dimensions");
    }
    if (!x.is_valid() || !F.is_valid() || !P.is_valid() || !Q.is_valid() || !R.is_valid()) {
        throw std::runtime_error("Invalid matrix state in predict");
    }
    Matrix<double> y = z - H * x;
    Matrix<double> S = H * P * H.transpose() + R;
    Matrix<double> K = P * H.transpose() * S.inverse();
    x = x + K * y;
    P = P - K * H * P;
}

Matrix<double> KalmanFilter::get_state() const {
    return x;
}

Matrix<double> KalmanFilter::get_state_ex() const {
    return x_ex;
} 